#include <set>
#include <limits>

template<typename D_>
KMean<D_>::KMean(boost::program_options::variables_map& vm, data_type& data)
    : vm_(vm), data_(data)
{
    K_ = vm.count("k") ? vm["k"].as<int>() : 2;
    maxiter_ = vm.count("maxiter") ? vm["maxiter"].as<int>() : 2;
}

template<typename D_>
void KMean<D_>::run()
{
    int n = data_.size();
    TextResult *output, *debug;

    output = !vm_.count("output") ? new TextResult(std::cout) : new TextResult(vm_["output"].as<std::string>());
    debug = !vm_.count("debug") ? 
        (!vm_.count("output") ? new TextResult(std::cout) : new TextResult(vm_["output"].as<std::string>()+"_debug")) 
        : new TextResult(vm_["debug"].as<std::string>());

    // init centers from the data points
    std::set<int> picked;
    for (int c = 0; c < K_; c++) {
        int r = rand() % n;
        while (ct(picked, r)) r = rand() % n;
        picked.insert(r);
        c_.pb(data_.getX(r));
    }

    // main loop
    std::vector<datum_type> ACC_center = c_;
    std::vector<int> ACC_count(K_);

    double total_dist = 1e-6;
    for (int iter = 0; iter < maxiter_; iter++) {
        // reset accumulation variables
        for (int c = 0; c < K_; c++) {
            ACC_center.at(c).zeros();
            ACC_count.at(c) = 0;
        }

        // assign points to closest centers
        double old_total = total_dist;
        total_dist = 0;
        for (int i = 0; i < n; i++) {
            const datum_type& x_i = data_.getX(i);
            int min_c = -1;
            double min_dist = DOUBLE_INFINITY;
            for (int c = 0; c < K_; c++) {
                double dist = arma::norm(c_.at(c) - x_i, 2);
                if (dist < min_dist) {
                    min_c = c;
                    min_dist = dist;
                }
            }
            total_dist += min_dist;
            ACC_center.at(min_c) += x_i;
            ACC_count.at(min_c) ++;
        }

        //MSG("iter = " << iter << " total_dist = " << total_dist);
        (*debug)() << iter << " " << total_dist; debug->flush();
        for (int c = 0; c < K_; c++) {
            (*debug)() << c_.at(c); debug->flush();
        }

        // check termination conditions
        //MSG(total_dist << " " << old_total << " c = " << fabs((old_total - total_dist) / old_total));
        if (fabs((old_total - total_dist) / old_total) < 1e-2) {
            break;
        }

        // update centers
        for (int c = 0; c < K_; c++) {
            if (ACC_count.at(c) != 0)
                c_.at(c) = ACC_center.at(c) / ACC_count.at(c);
            else { // no point assign to cluster c
                int r = rand() % n;
                c_.at(c) = data_.getX(r);
            }
        }

    }

    for (int i = 0; i < sz(c_); i++) {
        (*output)() << c_.at(i);
        output->flush();
    }

    delete output;
    delete debug;
}
